from PIL import Image
import torch
from torch.utils.data import Dataset
import pandas as pd
import numpy as np




class UnlabeledImageDataset(Dataset):
    def __init__(self, file_paths, transform=None, data_dir=''):
        self.file_paths = file_paths
        self.transform = transform
        self.data_dir = data_dir

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img_path = self.file_paths[idx]
        image = Image.open(self.data_dir + img_path)
        if self.transform is not None:
            image = self.transform(image)       
        patch_name = img_path.split('/')[-1].split('.')[0]
        
        return image, patch_name
    

class UnlabeledImageDatasetContrastive(Dataset):
    def __init__(self, file_paths, transform=None, data_dir=''):
        self.file_paths = file_paths
        self.transform = transform
        self.data_dir = data_dir

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img_path = self.file_paths[idx]
        image = Image.open(self.data_dir + img_path)

        if len(self.transform)>1:
            views = []
            for i in range(len(self.transform)):
                views.append(self.transform[i](image))
                
            image = views#view1, view2   
        elif self.transform is not None:
            image = self.transform[0](image)  
        patch_name = img_path.split('/')[-1].split('.')[0]
        
        return image, patch_name
    


class UCMercedMLDataset(Dataset):
    def __init__(self, file_paths, transform=None, data_dir='', multi_label_dict='multi-labels.xlsx'):

        self.file_paths = file_paths
        self.transform = transform
        self.multi_label_dict = pd.read_excel(multi_label_dict, sheet_name='Sheet1').set_index(
            'IMAGE\LABEL').T.to_dict('list')
        self.data_dir = data_dir


    def __len__(self):
        return len(self.file_paths)

    
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img_path = self.file_paths[idx]
        image = Image.open(self.data_dir + img_path)       

        image_id = img_path.split('/')[-1].split('.')[0]
        label = np.array(self.multi_label_dict[image_id])

        if self.transform is not None:
            image = self.transform(image) 
        
        return image, label, image_id#patch_name
    

DLRS_lookup = {
    'airplane': {'id':0, 'val': [ 166, 202, 240 ]},
    'baresoil': {'id':1, 'val':  [ 128, 128, 0 ]},
    'buildings': {'id':2, 'val':  [ 0, 0, 128 ]},
    'cars': {'id':3, 'val':  [ 255, 0, 0 ]},
    'chaparral': {'id':4, 'val':  [ 0, 128, 0 ]},
    'court': {'id':5, 'val':  [ 128, 0, 0 ]},
    'dock': {'id':6, 'val':  [ 255, 233, 233 ]},
    'field': {'id':7, 'val':  [ 160, 160, 164 ]},
    'grass': {'id':8, 'val':  [ 0, 128, 128]},
    'mobilehome': {'id':9, 'val':  [ 90, 87, 255]},
    'pavement': {'id':10, 'val':  [ 255, 255, 0 ]},
    'sand': {'id':11, 'val':  [ 255, 192, 0 ]},
    'sea': {'id':12, 'val':  [ 0, 0, 255, ]},
    'ship': {'id':13, 'val':  [ 255, 0 , 192 ]},
    'tanks': {'id':14, 'val':  [ 128, 0, 128 ]},
    'trees': {'id':15, 'val':  [ 0, 255, 0 ]},
    'water': {'id':16, 'val':  [ 0, 255, 255 ]}
}

DLRS_lookup_reverse = {
    0:'airplane',
    1:'baresoil',
    2:'buildings',
    3:'cars',
    4:'chaparral',
    5:'court',
    6:'dock',
    7:'field',
    8:'grass',
    9:'mobilehome',
    10:'pavement',
    11:'sand',
    12:'sea',
    13:'ship',
    14:'tanks',
    15:'trees',
    16:'water'
}
